#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@author: liang kang
@contact: gangkanli1219@gmail.com
@time: 1/3/18 5:18 PM
@desc: trying to speed up predicting
"""
import logging
import os
import time

import cv2
import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
from sklearn.metrics import precision_score
from sklearn.metrics import recall_score

from .predict_low_density_image import get_center_boxes_classes_scores
from .predict_low_density_image import get_edge_boxes_classes_scores
from .utils import visualization_utils as vis_utils
from .utils.basic import get_file_name
from .utils.eval import eval_detect_result
from .utils.eval import load_model
from .utils.eval import merge_region_prediction
from .utils.eval import read_xml_as_eval_info


def predict_image(image_path, output=None,
                  checkpoint_center='', checkpoint_edge='',
                  score=0.3, percent=0.9):
    """
    预测函数，并统计

    Parameters
    ----------
    image_path: 图片文件夹
    output: 保存预测和的图像，为空则不保存
    checkpoint_center: 中心模型
    checkpoint_edge: 边缘模型
    score: 显示的置信度阈值
    percent: 矩形框重叠率阈值

    Returns
    -------

    """
    start_time = time.time()
    detection_graph_center = load_model(checkpoint_center)
    detection_graph_edge = load_model(checkpoint_edge)
    config = tf.ConfigProto(device_count={"CPU": 4, "GPU": 1})
    config.gpu_options.allow_growth = True
    stage_time = time.time()
    logging.info('{} elapsed time: {:.3f}s'.format(time.ctime(),
                                                   stage_time - start_time))
    logging.info('loading models completed !')
    start_time = stage_time
    ground_true = []
    predictions = []
    all_pic_recall = []
    all_pic_pre = []
    for idx, image_name in enumerate(os.listdir(image_path)):
        if image_name.split('.')[-1] == 'png' or image_name.split('.')[-1] == 'jpg':
            logging.info('predicting image: {}!'.format(os.path.join(image_path, image_name)))
            img_head_name = get_file_name(image_name)
            img = cv2.imread(os.path.join(image_path, image_name))
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            xml_path = os.path.join(image_path, img_head_name + '.xml')

            _boxes, _classes, _scores = [], [], []
            get_center_boxes_classes_scores(
                _boxes, _classes, _scores, config, detection_graph_center, img)
            get_edge_boxes_classes_scores(
                _boxes, _classes, _scores, config, detection_graph_edge, img)

            stage_time = time.time()
            logging.info('{} elapsed time: {:.3f}s'.format(time.ctime(),
                                                           stage_time - start_time))
            logging.info('predicting image completed !')
            start_time = stage_time

            boxes, classes, scores = [], [], []
            for box, cls, sc in zip(_boxes, _classes, _scores):
                if sc >= score:
                    boxes.append(box)
                    classes.append(cls)
                    scores.append(sc)
            _boxes, _classes, _scores = merge_region_prediction(
                np.array(boxes), np.array(scores), np.array(classes), percent)
            _boxes = np.array(_boxes)
            _classes = np.array(_classes).astype(np.int32)
            _scores = np.array(_scores)
            objects = read_xml_as_eval_info(xml_path, ['person'])['objects']
            gt_classes, gt_boxes = np.hsplit(np.array(objects), [1])
            if len(objects) == 0:
                gt_boxes = []
                gt_classes = []
            else:
                gt_classes = list(gt_classes.flatten())
                gt_boxes = gt_boxes[:, (1, 0, 3, 2)]
                gt_boxes = list(map(lambda x: list(x), gt_boxes))
            if output is not None:
                vis_utils.visualize_boxes_and_labels_on_image_array(
                    img, np.array(_boxes), np.array(_classes).astype(np.uint8), np.array(_scores),
                    use_normalized_coordinates=True,
                    min_score_thresh=score,
                    line_thickness=3,
                    max_boxes_to_draw=200)
                plt.figure(figsize=(10, 10), dpi=250)
                plt.title(img_head_name)
                plt.imshow(img)
                plt.savefig(os.path.join(output, '{}.jpg'.format(idx)))
                plt.close()
            gt, pred = eval_detect_result(gt_boxes, list(map(lambda x: x + 1, gt_classes)),
                                          _boxes, _classes, default_class=0)
            ground_true += gt
            predictions += pred
            recall = recall_score(gt, pred)
            precision = precision_score(gt, pred)
            all_pic_recall.append(recall)
            all_pic_pre.append(precision)
            stage_time = time.time()
            logging.info('{} elapsed time: {:.3f}s'.format(time.ctime(),
                                                           stage_time - start_time))
            logging.info('computing statistics completed !')
            start_time = stage_time

    recall = recall_score(ground_true, predictions)
    precision = precision_score(ground_true, predictions)
    logging.info('Total Recall: {:.5f}'.format(recall))
    logging.info('Total Precision: {:.5f}'.format(precision))
    logging.info('Mean Recall: {:.5f}'.format(np.mean(all_pic_recall)))
    logging.info('Mean Precision: {:.5f}'.format(np.mean(all_pic_pre)))


def predict_with_recall_precision(root, output=None, score=0.5, percent=0.7):
    """
    预测函数，并统计

    Parameters
    ----------
    root: 图片文件夹
    output: 保存预测和的图像，为空则不保存
    score: 显示的置信度阈值
    percent: 矩形框重叠率阈值

    Returns
    -------

    """
    logging.basicConfig(level=logging.INFO)
    if (output is not None) and (not os.path.exists(output)):
        os.makedirs(output)
    # 中心模型的路径
    checkpoint_center = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                                     'data', 'low', 'frozen_inference_graph_fast_center_200000.pb')
    # 边缘模型的路径
    checkpoint_edge = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                                   'data', 'low', 'frozen_inference_graph_fast_edge_160998.pb')
    predict_image(
        image_path=root,
        output=output,
        checkpoint_center=checkpoint_center,
        checkpoint_edge=checkpoint_edge,
        score=score,
        percent=percent)
